from typing import Union, Dict, Any
from fastapi import APIRouter, HTTPException, BackgroundTasks
from app.services.dhcds_api_service import DynamicHeterogeneousCommunityDetection
from app.schemas.dhgam.dhcds_base_schema import AlgorithmRequest, AlgorithmResponse, AlgorithmMiddleResponse, TaskStatus
from app.tools.redis_client import redis_client
from app.utils.logger import logger
import networkx as nx
import numpy as np

router = APIRouter()
service = DynamicHeterogeneousCommunityDetection()

@router.post("/create", response_model=AlgorithmMiddleResponse)
async def create_task(request: AlgorithmRequest) -> AlgorithmMiddleResponse:
    """
    创建新的动态异构图社区检测(DHCDS)算法任务并存储在Redis中
    
    Args:
        request (AlgorithmRequest): 包含任务ID、回调URL和输入参数的请求对象
        
    Returns:
        AlgorithmMiddleResponse: 包含任务初始状态的响应对象
        
    Raises:
        HTTPException: 当任务创建失败时抛出500错误
    """
    try:
        task_response = AlgorithmMiddleResponse(
            task_id=request.task_id,
            task_callback_url=request.task_callback_url,
            task_status=TaskStatus.PENDING,
            task_progress=0,
            input_params=request.input_params,
            metrics=[]
        )
        await redis_client.set_data(request.task_id, task_response.model_dump())
        return task_response
    except Exception as e:
        logger.error(f"创建DHCDS任务失败: {str(e)}")
        raise HTTPException(status_code=500, detail=f"创建DHCDS任务失败: {str(e)}")

async def add_task_to_background(
    background_tasks: BackgroundTasks,
    input_params: Dict[str, Any],
    task_id: str
) -> None:
    """
    将DHCDS任务添加到后台执行队列
    
    Args:
        background_tasks: FastAPI的后台任务对象
        input_params: 算法输入参数
        task_id: 任务唯一标识
    """
    background_tasks.add_task(run_task_in_background, input_params, task_id)

@router.post("/execute/{task_id}", response_model=AlgorithmMiddleResponse)
async def execute_task(
    task_id: str, 
    background_tasks: BackgroundTasks
) -> AlgorithmMiddleResponse:
    """
    执行指定ID的DHCDS任务
    
    Args:
        task_id: 任务唯一标识
        background_tasks: FastAPI的后台任务对象
        
    Returns:
        AlgorithmMiddleResponse: 任务执行状态的响应对象
        
    Raises:
        HTTPException: 当任务不存在时抛出404错误，执行失败时抛出500错误
    """
    try:
        task_data = await redis_client.get_data(task_id)
        if not task_data:
            raise HTTPException(status_code=404, detail="DHCDS任务不存在")

        input_params = task_data.get("input_params")
        result_key = f'{task_id}_result'
        
        # 更新任务状态为执行中
        running_status = {
            "task_id": task_id,
            "task_callback_url": task_data.get("task_callback_url"),
            "task_status": TaskStatus.RUNNING,
            "task_progress": 0,
            "input_params": input_params,
            "metrics": []
        }
        await redis_client.set_data(result_key, running_status)
        
        # 添加到后台任务队列
        await add_task_to_background(background_tasks, input_params, task_id)
        
        return AlgorithmMiddleResponse(**running_status)
    except HTTPException:
        raise
    except Exception as e:
        logger.error(f"执行DHCDS任务失败: {str(e)}")
        raise HTTPException(status_code=500, detail=f"执行DHCDS任务失败: {str(e)}")

@router.get("/result/{task_id}", response_model=Union[AlgorithmResponse, AlgorithmMiddleResponse])
async def get_result(task_id: str) -> Union[AlgorithmResponse, AlgorithmMiddleResponse]:
    """
    获取DHCDS任务执行结果
    
    Args:
        task_id: 任务唯一标识
        
    Returns:
        Union[AlgorithmResponse, AlgorithmMiddleResponse]: 
            - 如果任务完成，返回AlgorithmResponse
            - 如果任务进行中，返回AlgorithmMiddleResponse
            
    Raises:
        HTTPException: 当结果不存在时抛出404错误，获取失败时抛出500错误
    """
    try:
        result_key = f'{task_id}_result'
        result_data = await redis_client.get_data(result_key)
        if not result_data:
            raise HTTPException(status_code=404, detail="DHCDS结果不存在")
        
        # 检查并修正社区演化数据格式
        if "output_params" in result_data and "dhcds_results" in result_data["output_params"]:
            dhcds_results = result_data["output_params"]["dhcds_results"]
            if "community_evolution" in dhcds_results and isinstance(dhcds_results["community_evolution"], dict):
                # 将字典转换为列表格式
                community_evolution_list = []
                for comm_id, members in dhcds_results["community_evolution"].items():
                    community_evolution_list.append({
                        "community_id": comm_id,
                        "members": members
                    })
                dhcds_results["community_evolution"] = community_evolution_list
        
        return (AlgorithmResponse(**result_data) 
                if "output_params" in result_data 
                else AlgorithmMiddleResponse(**result_data))
    except HTTPException:
        raise
    except Exception as e:
        logger.error(f"获取DHCDS结果失败: {str(e)}")
        raise HTTPException(status_code=500, detail=f"获取DHCDS结果失败: {str(e)}")

def convert_sets_to_lists(obj):
    """
    递归地将对象中的所有集合和NumPy数组转换为Python原生类型
    
    Args:
        obj: 任意 Python 对象
        
    Returns:
        转换后的对象（所有集合变为列表，NumPy数组变为列表）
    """
    if isinstance(obj, set):
        return list(obj)
    elif isinstance(obj, np.ndarray):
        return obj.tolist()
    elif isinstance(obj, np.integer):
        return int(obj)
    elif isinstance(obj, np.floating):
        return float(obj)
    elif isinstance(obj, dict):
        return {k: convert_sets_to_lists(v) for k, v in obj.items()}
    elif isinstance(obj, list):
        return [convert_sets_to_lists(v) for v in obj]
    else:
        return obj

async def run_task_in_background(input_params: Dict[str, Any], task_id: str) -> None:
    """
    在后台执行DHCDS算法任务
    
    Args:
        input_params: 算法输入参数
        task_id: 任务唯一标识
    """
    try:
        # 创建服务实例并配置参数
        dhcds_params = input_params.get("dhcds_params", {})
        service = DynamicHeterogeneousCommunityDetection(
            stability_weight=dhcds_params.get("stability_weight", 0.5),
            embedding_dim=dhcds_params.get("embedding_dim", 64),
            num_clusters=dhcds_params.get("num_clusters", 10)
        )
        
        # 准备图数据
        graph_sequence = []
        for graph_data in dhcds_params.get("graph_sequence", []):
            G = nx.Graph()
            
            # 添加节点
            for node in graph_data.get("nodes", []):
                G.add_node(node["id"], type=node["type"], **node.get("attributes", {}))
            
            # 添加边
            for edge in graph_data.get("edges", []):
                G.add_edge(
                    edge["source"], 
                    edge["target"], 
                    type=edge["type"],
                    weight=edge["weight"],
                    **edge.get("attributes", {})
                )
            
            graph_sequence.append(G)
        logger.info(f"图序列: {graph_sequence}")
        # 设置节点类型权重（如果提供）
        if dhcds_params.get("node_type_weights"):
            service.node_type_weights = dhcds_params.get("node_type_weights")
        
        # 执行社区检测
        results = service.detect_dynamic_communities(graph_sequence)
        
        # 准备结果
        output_results = []
        for i, result in enumerate(results):
            # 确保所有集合类型转换为列表
            communities = {k: list(v) if isinstance(v, set) else v for k, v in result["communities"].items()} if isinstance(result["communities"], dict) else result["communities"]
            type_communities = {k: list(v) if isinstance(v, set) else v for k, v in result["type_communities"].items()} if isinstance(result["type_communities"], dict) else result["type_communities"]
            evolution = {k: list(v) if isinstance(v, set) else v for k, v in result["evolution"].items()} if isinstance(result["evolution"], dict) else result["evolution"]
            
            output_results.append({
                "time": i,
                "communities": communities,
                "type_communities": type_communities,
                "evolution": evolution,
                "community_embeddings": result["community_embeddings"]
            })
        
        # 确保 community_evolution 是列表格式
        community_evolution_list = []
        if hasattr(service, 'community_evolution') and service.community_evolution:
            logger.info(f"社区演化数据类型: {type(service.community_evolution)}")
            
            for k, v in service.community_evolution.items():
                community_evolution_list.append({
                    "community_id": k,
                    "members": list(v) if isinstance(v, set) else v
                })

        output_params = {
            "dhcds_results": {
                "results": output_results,
                "algorithm": "DHCDS",
                "parameters": {
                    "stability_weight": service.stability_weight,
                    "embedding_dim": service.embedding_dim,
                    "num_clusters": service.num_clusters,
                    "node_type_weights": service.node_type_weights
                },
                "community_evolution": community_evolution_list  # 使用列表格式
            }
        }
        
        # 创建最终响应并在存储前转换所有不可序列化对象
        final_result = convert_sets_to_lists({
            "task_id": task_id,
            "task_status": TaskStatus.COMPLETED,
            "task_progress": 100,
            "output_params": output_params,
            "metrics": []
        })
        
        await redis_client.update_data(f'{task_id}_result', final_result)
    except Exception as e:
        logger.error(f"DHCDS后台任务执行失败: {str(e)}")
        # 更新任务状态为失败
        error_status = {
            "task_id": task_id,
            "task_status": TaskStatus.FAILED,
            "error_message": str(e)
        }
        await redis_client.update_data(f'{task_id}_result', error_status) 